import torch
import torch.nn as nn


class DiscriminatorLSTMFransTrue(nn.Module):
    def __init__(self, input_dim, hidden_dim, out_dim):
        super(DiscriminatorLSTMFransTrue, self).__init__()
        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers=2, bidirectional=True)
        self.out = nn.Linear(hidden_dim*2, out_dim)
        # pass
    # def forward(self, input, lyrics):
    #     concat_input = torch.cat((input, lyrics), 2)
    #     # print(concat_input)
    #     # print(concat_input.shape)
    #     _, lstm_out = self.lstm(concat_input.view(concat_input.shape[1], concat_input.shape[0], -1))
    #     # lstm_out two tuple, each (2,100,400)
    #
    #     ct = lstm_out[0]
    #
    #     # print(ct[-1])
    #     last_layer_out = ct[-1]
    #
    #
    #     last_layer_out = torch.unsqueeze(last_layer_out,dim=1)
    #
    #
    #     out_line=self.out(last_layer_out)
    #     out = torch.sigmoid(out_line)
    #
    #     # out = self.out(last_layer_out)
    #     out = out.view(out.shape[0], -1)
    #
    #     return out

    def forward(self, input, lyrics):
        concat_input = torch.cat((input, lyrics), 2)
        lstm_out, _ = self.lstm(concat_input)
        out_line = self.out(lstm_out)
        out_line = torch.sigmoid(out_line)
        out_line=torch.mean(out_line,1)
        return out_line
